package com.example.endusation.ollama;

import com.example.endusation.dto.OllamaChatResponseDTO;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.converter.StringHttpMessageConverter;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

@Component
public class OllamaDirectClient {

    private final String ollamaApiUrl;
    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;

    // 从配置文件读取Ollama服务地址
    public OllamaDirectClient(@Value("${ollama.host:localhost}") String host,
                              @Value("${ollama.port:11434}") int port) {
        this.ollamaApiUrl = "http://" + host + ":" + port + "/api/chat";
        this.restTemplate = new RestTemplate();
        // 1. 移除默认的 StringHttpMessageConverter（默认用 ISO-8859-1）
        this.restTemplate.getMessageConverters().removeIf(converter ->
                converter instanceof StringHttpMessageConverter
        );

// 2. 添加 UTF-8 编码的 StringHttpMessageConverter
        this.restTemplate.getMessageConverters().add(new StringHttpMessageConverter(StandardCharsets.UTF_8));
        this.objectMapper = new ObjectMapper();
    }

    /**
     * 调用Ollama模型生成回答
     */
    public String chat(String model, String prompt) {
        try {
            // 构建请求参数
            Map<String, Object> request = new HashMap<>();
            request.put("model", model);

            // 构建消息列表
            Map<String, String> message = new HashMap<>();
            message.put("role", "user");
            message.put("content", prompt);
            request.put("messages", new Object[]{message});
            request.put("stream", false);

            // 发送请求
            OllamaChatResponseDTO response = restTemplate.postForObject(ollamaApiUrl, request, OllamaChatResponseDTO.class);

            // 解析响应
            return response.getMessage().getContent();
        } catch (Exception e) {
            throw new RuntimeException("调用Ollama API失败: " + e.getMessage(), e);
        }
    }
}
